# -*- coding: utf-8 -*-
# lio_transfer 是把lio获得的位置传递给飞控

import rospy
from geometry_msgs.msg import PoseStamped, Point
from nav_msgs.msg import Odometry
from sensor_msgs.msg import PointCloud2
import sensor_msgs.point_cloud2 as pc2
import math
from pyquaternion import Quaternion
import tf

# BIAS = [[0.0, 0.0, 0.0],
#     [3.0, 0.0, 0.0],
#     [-3.0, 0.0, 0.0],
#     [-3.0, -3.0, 0.0]]
BIAS = [0.0, 0.0, 0.0]
uav_id = 0 
local_poses = None  # 存储每个无人机的 PoseStamped 消息
quaternions = None  # 存储每个无人机的初始旋转四元数

# 初始化每个无人机的 PoseStamped 消息和初始四元数

local_poses = PoseStamped()
local_poses.header.frame_id = 'world'
quaternions = tf.transformations.quaternion_from_euler(0, 0, 0)

# 回调函数处理每架无人机的 Odometry 数据
def lio_callback(data):
    # 更新位置
    local_poses.pose.position.x = data.pose.pose.position.x + BIAS[0]
    local_poses.pose.position.y = data.pose.pose.position.y + BIAS[1]
    local_poses.pose.position.z = data.pose.pose.position.z + BIAS[2]

    # 更新旋转
    q_ = Quaternion([data.pose.pose.orientation.w, data.pose.pose.orientation.x, data.pose.pose.orientation.y, data.pose.pose.orientation.z])
    q = Quaternion([quaternions[3], quaternions[0], quaternions[1], quaternions[2]])
    q_ = q_ * q
    local_poses.pose.orientation.w = q_[0]
    local_poses.pose.orientation.x = q_[1]
    local_poses.pose.orientation.y = q_[2]
    local_poses.pose.orientation.z = q_[3]

    # 发布偏移后的 Odometry
    transfer_odometry = data
    transfer_odometry.pose.pose.position.x = data.pose.pose.position.x + BIAS[0]
    transfer_odometry.pose.pose.position.y = data.pose.pose.position.y + BIAS[1]
    transfer_odometry.pose.pose.position.z = data.pose.pose.position.z + BIAS[2]
    odom_pubs.publish(transfer_odometry)

# 回调函数处理每架无人机的 PointCloud2 数据
def cloud_callback(data):
    # 提取点云所有字段的数据
    cloud_points = pc2.read_points(
        data, 
        skip_nans=True, 
        field_names=("x", "y", "z", "intensity", "normal_x", "normal_y", "normal_z", "curvature")
    )
    
    new_points = []
    for point in cloud_points:
        # 更新 x, y, z 坐标，加上偏移
        new_points.append((
            # point[0],  # x 坐标
            # point[1],  # y 坐标
            # point[2],  # z 坐标
            point[0] + BIAS[0],  # x 坐标
            point[1] + BIAS[1],  # y 坐标
            point[2] + BIAS[2],  # z 坐标
            point[3],                      # intensity 保持不变
            point[4],                      # normal_x 保持不变
            point[5],                      # normal_y 保持不变
            point[6],                      # normal_z 保持不变
            point[7]                       # curvature 保持不变
        ))

    # 定义点云的字段
    fields = [
        pc2.PointField('x', 0, pc2.PointField.FLOAT32, 1),
        pc2.PointField('y', 4, pc2.PointField.FLOAT32, 1),
        pc2.PointField('z', 8, pc2.PointField.FLOAT32, 1),
        pc2.PointField('intensity', 32, pc2.PointField.FLOAT32, 1),
        pc2.PointField('normal_x', 16, pc2.PointField.FLOAT32, 1),
        pc2.PointField('normal_y', 20, pc2.PointField.FLOAT32, 1),
        pc2.PointField('normal_z', 24, pc2.PointField.FLOAT32, 1),
        pc2.PointField('curvature', 36, pc2.PointField.FLOAT32, 1)
    ]

    # 创建新的点云消息
    header = data.header
    header.frame_id = "world"  # 修改 frame_id（可选）
    new_cloud = pc2.create_cloud(header, fields, new_points)
    
    # 发布新点云
    cloud_pubs.publish(new_cloud)
    
    
    

if __name__ == '__main__':
    rospy.init_node('lio_transfer_'+str(uav_id))

    # 为每架无人机初始化对应的发布器

    position_pubs = rospy.Publisher("/iris_{}/mavros/vision_pose/pose".format(uav_id), PoseStamped, queue_size=1)
    cloud_pubs = rospy.Publisher("/iris_{}/transfer/cloud_registered".format(uav_id), PointCloud2, queue_size=1)
    odom_pubs = rospy.Publisher("/iris_{}/transfer/Odometry".format(uav_id), Odometry, queue_size=1)

    rospy.Subscriber("iris_{}/Odometry".format(uav_id), Odometry, lio_callback, queue_size=1)
    rospy.Subscriber("iris_{}/cloud_registered".format(uav_id), PointCloud2, cloud_callback, queue_size=1)

    rate = rospy.Rate(30)  # 发布频率 30Hz

    while not rospy.is_shutdown():

        if local_poses.pose.position == Point():
            continue
        else:
            local_poses.header.stamp = rospy.Time.now()
            position_pubs.publish(local_poses)

        try:
            rate.sleep()
        except:
            continue
